package cn.aezo.chat_gpt.modules.chat.websocket;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.gson.Gson;
import com.unfbx.chatgpt.entity.chat.ChatCompletionResponse;
import com.unfbx.chatgpt.entity.chat.Message;
import okhttp3.Response;
import okhttp3.WebSocket;
import okhttp3.WebSocketListener;
import okio.ByteString;

import javax.websocket.Session;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;

public class XfYunAiWebsoketListener extends WebSocketListener {

    private static ConcurrentHashMap<String, Session> sessionMap;
    private  String message;

    private String appId;

    private String uid;

    public static ConcurrentHashMap<String, List<String>> answerMap = new ConcurrentHashMap<>(); // 大模型的答案汇总

    public static final Gson gson = new Gson();


    public XfYunAiWebsoketListener(String msg,String appId, String uid, ConcurrentHashMap<String, Session> ssMap) {
        sessionMap = ssMap;
        this.message = msg;
        this.appId = appId;
        this.uid = uid;
        answerMap.put(uid,new ArrayList<String>());
    }

    @Override
    public void onOpen(WebSocket webSocket, Response response) {
        // 连接已打开
        System.out.println("连接已打开");

        JSONObject requestJson=new JSONObject();

        JSONObject header=new JSONObject();  // header参数
        header.put("app_id",this.appId);
        header.put("uid", UUID.randomUUID().toString().substring(0, 10));

        JSONObject parameter=new JSONObject(); // parameter参数
        JSONObject chat=new JSONObject();
        chat.put("domain","generalv2");
        chat.put("temperature",0.5);
        chat.put("max_tokens",4096);
        parameter.put("chat",chat);

        JSONObject payload=new JSONObject(); // payload参数
        JSONObject message=new JSONObject();
        JSONArray text=new JSONArray();

        // 历史问题获取
        Message msg = new Message();
        msg.setRole("user");
        msg.setContent(this.message);
        text.add(JSON.toJSON(msg));

        message.put("text",text);
        payload.put("message",message);


        requestJson.put("header",header);
        requestJson.put("parameter",parameter);
        requestJson.put("payload",payload);
        webSocket.send(requestJson.toString());
    }

    @Override
    public void onMessage(WebSocket webSocket, String text) {
        // 收到文本消息
        // System.out.println(userId + "用来区分那个用户的结果" + text);
        JsonParse myJsonParse = gson.fromJson(text, JsonParse.class);
        if (myJsonParse.header.code != 0) {
            System.out.println("发生错误，错误码为：" + myJsonParse.header.code);
            System.out.println("本次请求的sid为：" + myJsonParse.header.sid);
            sendMsgByUid(this.uid,"sqchat","服务异常，错误码："+ myJsonParse.header.code);
            webSocket.close(1000, "");
            return;
        }
        List<Text> textList = myJsonParse.payload.choices.text;

        for (Text value : textList) {
            sendMsgByUid(this.uid, answerMap.get(this.uid).isEmpty() ? "assistant" : "", value.content);
            answerMap.get(this.uid).add(value.content);
        }
    }

    @Override
    public void onMessage(WebSocket webSocket, ByteString bytes) {
        // 收到二进制消息
    }

    @Override
    public void onClosing(WebSocket webSocket, int code, String reason) {
        // 连接正在关闭
        System.out.println("连接正在关闭"+this.uid);
        answerMap.put(this.uid,new ArrayList<>());
    }

    @Override
    public void onClosed(WebSocket webSocket, int code, String reason) {
        // 连接已关闭
        System.out.println("连接已关闭"+this.uid);
    }

    @Override
    public void onFailure(WebSocket webSocket, Throwable t, Response response) {
        // 连接失败
        answerMap.put(this.uid,new ArrayList<>());
        System.out.println("连接失败"+this.uid);
        sendMsgByUid(this.uid,"sqchat","服务器开小差了，请重试~");
    }

    public static void sendMsgByUid(String uid,String role,String content) {
        Message roleContent=new Message();
        roleContent.setRole(role);
        roleContent.setContent(content);
        ObjectMapper mapper = new ObjectMapper();
        try {
            String delta  = mapper.writeValueAsString(roleContent);
            sessionMap.get(uid).getBasicRemote().sendText(delta);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    private class RoleContent{
        String role;
        String content;

        public String getRole() {
            return role;
        }

        public void setRole(String role) {
            this.role = role;
        }

        public String getContent() {
            return content;
        }

        public void setContent(String content) {
            this.content = content;
        }
    }

    //返回的json结果拆解
    private class JsonParse {
        Header header;
        Payload payload;
    }

    private class Header {
        int code;
        int status;
        String sid;
    }

    private class Payload {
        Choices choices;
    }

    private class Choices {
        List<Text> text;
    }

    private class Text {
        String role;
        String content;
    }

}
